import glob

from inner_utils import *
from segment_airway import segment_airway
from segment_lung import segment_lung

params = define_parameter()


def generate_mask_for_train():
    save_root_path = r"D:\dataset\vessel\Parse2022_valid_data"
    train_data_root_path = r"E:\luanqibazao\evaluation_2022\evaluation\*"
    sample_list = glob.glob(train_data_root_path)
    sample_list = sample_list
    for sample_path in sample_list:
        # sample_path = r"D:\dataset\vessel\Parse_2022_train_data\PA000309"
        #####################################################
        # Load image
        #####################################################
        segment_result = glob.glob(os.path.join(sample_path, "image", "PA*_lung.nii.gz"))
        if len(segment_result) > 0:
            print(segment_result[0] + " has existed!")
            continue
        file_path = os.path.join(sample_path, "image", "PA*.nii.gz")
        file_list = glob.glob(file_path)
        file_path = file_list[0]
        I = nib.load(file_path)
        params['currentSampleRootPath'] = os.path.dirname(file_path)
        params['currentSampleName'] = os.path.basename(file_path).split(".")[0]
        print("start to load " + params["currentSampleName"])
        I_affine = I.affine
        I = I.get_data()

        #####################################################
        # Coarse segmentation of lung & airway
        #####################################################

        Mlung = segment_lung(params, I, I_affine)
        print("The lung of " + params["currentSampleName"] + " has been segmented")

        #####################################################
        # Romove airway from lung mask
        #####################################################

        Mlung, Maw = segment_airway(params, I, I_affine, Mlung)
        print("The airway of " + params["currentSampleName"] + " has been segmented")


def generate_mask_for_valid():
    save_root_path = r"C:\dataset\vessel\Parse2022_valid_data"
    train_data_root_path = r"C:\dataset\vessel\Parse2022_valid_data\*"
    sample_list = glob.glob(train_data_root_path)
    sample_list = sample_list
    for sample_path in sample_list:
        sample_name = os.path.basename(sample_path).split(".")[0]
        # sample_path = r"D:\dataset\vessel\Parse_2022_train_data\PA000309"
        #####################################################
        # Load image
        #####################################################
        segment_result = glob.glob(os.path.join(sample_path, "image", "PA*_lung.nii.gz"))
        if len(segment_result) > 0:
            print(segment_result[0] + " has existed!")
            continue
        save_path = os.path.join(save_root_path,sample_name)
        if not os.path.exists(save_path):
            os.mkdir(save_path)
        save_path = os.path.join(save_path, "image")
        if not os.path.exists(save_path):
            os.mkdir(save_path)

        I = nib.load(sample_path)
        I.to_filename(os.path.join(save_path,sample_name+".nii.gz"))
        params['currentSampleRootPath'] = save_path
        params['currentSampleName'] = sample_name
        print("start to load " + params["currentSampleName"])
        I_affine = I.affine
        I = I.get_data()

        #####################################################
        # Coarse segmentation of lung & airway
        #####################################################

        Mlung = segment_lung(params, I, I_affine)
        print("The lung of " + params["currentSampleName"] + " has been segmented")

        #####################################################
        # Romove airway from lung mask
        #####################################################

        Mlung, Maw = segment_airway(params, I, I_affine, Mlung)
        print("The airway of " + params["currentSampleName"] + " has been segmented")


if __name__ == '__main__':
    generate_mask_for_valid()